def predict(model, window_rgb):
    results = model.predict(window_rgb, conf=0.2)
    return data_handle(results)


def data_handle(results):
    """
      0: me
      1: enemy
      2: pickup
      3: isfight
      4: direction
      5: door
      6: card
      7: challenge_again
      8: back_town
    """
    names = results[0].names
    boxes = results[0].boxes
    cls = boxes.cls.cpu().numpy()
    locs = boxes.xyxy.cpu().numpy()
    # conf = boxes.conf.cpu().numpy()
    # print(f"conf: {conf}")

    data_dic = dict()
    for _index, _name in names.items():
        _name_data = []

        for index_cls, _name_cls in enumerate(cls):
            na = names[_name_cls]
            if na == _name:
                loc = locs[index_cls]
                left_x = loc[0]
                left_y = loc[1]
                right_x = loc[2]
                right_y = loc[3]

                _name_data.append((left_x, left_y, right_x, right_y))
        if len(_name_data) > 0:
            data_dic[_name] = _name_data
    return data_dic


def get_down_mid_loc(left_x, left_y, right_x, right_y):
    return (round((left_x + right_x) / 2), round(right_y))


def get_all_mid_loc(left_x, left_y, right_x, right_y):
    return (round((left_x + right_x) / 2), round((left_y + right_y) / 2))


def get_data_all(data_dic, key, is_all_mid=True):
    results = []
    locs = data_dic.get(key)
    if locs is not None and len(locs) > 0:
        for index, loc in enumerate(locs):
            if is_all_mid:
                results.append(get_all_mid_loc(loc[0], loc[1], loc[2], loc[3]))
            else:
                results.append(get_down_mid_loc(loc[0], loc[1], loc[2], loc[3]))
    return results


def get_data_one(data_dic, key, is_all_mid=True):
    locs = data_dic.get(key)
    if locs is not None and len(locs) > 0:
        loc = locs[0]
        if is_all_mid:
            return True, get_all_mid_loc(loc[0], loc[1], loc[2], loc[3])
        else:
            return True, get_down_mid_loc(loc[0], loc[1], loc[2], loc[3])
    return False, (0, 0)
